package lab.partitioner

import org.apache.spark.{HashPartitioner, Partitioner, SparkConf, SparkContext}

/**
  * Created by kawhi on 30/03/2017.
  * Edited by Q.
  */

object PartialKey {
  def showPartition[T](i: Int, iter: Iterator[T]) : Iterator[(T)] = {
    iter.foreach(ele => {
      println("partition\t" + i + "\t" + ele)
    })
    iter
  }

  def main(args: Array[String]) {
    val conf = new SparkConf().setAppName("Partial Key Partitioner").setMaster("local[4]")
    val sc = new SparkContext(conf)
    val ports = sc.parallelize(List('a','b','c', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a'),4).map((_, 1))
//    print(ports.getNumPartitions)
//    将HashPartitioner换成 PartialKeyPartitioner
    val wbt = ports.partitionBy(new PartialKeyPartitioner(3))/*.reduceByKey((a, b)=>a+b)*/
    val res = wbt.mapPartitionsWithIndex(showPartition).collect

    sc.stop()
  }
}


class PartialKeyPartitioner(partitions: Int) extends Partitioner {
  def numPartitions : scala.Int = {
    partitions
  }
  //val partitionBucket = new PartitionBucket
  object partitionBucket {
    var buckets = new Array[Int](partitions)
    def get(position: Int): Int = {
      buckets(position)
    }
    def plus(position: Int) {
      buckets(position) = buckets(position) + 1
    }
  }
  def getPartition(key: scala.Any): scala.Int = {
    val partition_1 = (key.hashCode() & Integer.MAX_VALUE) % partitions
    val partition_2 = ((key.hashCode() & Integer.MAX_VALUE) + 1) % partitions

//    println("before sending:")
//    println("key: " + key + " two choice: " + partition_1.toString + " or " + partition_2.toString)
//    println(partitionBucket.get(0) + " " + partitionBucket.get(1) + " " + partitionBucket.get(2) + " ")

    if (partitionBucket.get(partition_1) < partitionBucket.get(partition_2)) {
      partitionBucket.plus(partition_1)
//      println("after sending:")
//      println(partitionBucket.get(0) + " " + partitionBucket.get(1) + " " + partitionBucket.get(2) + " ")
//      println()
      return partition_1
    }
    else {
      partitionBucket.plus(partition_2)
//      println(partitionBucket.get(0) + " " + partitionBucket.get(1) + " " + partitionBucket.get(2) + " ")
//      println("after sending:")
//      println(partitionBucket.get(0) + " " + partitionBucket.get(1) + " " + partitionBucket.get(2) + " ")
//      println()
      return partition_2
    }

  }
}



